import numpy as np
import torch
from DDPM.Network import DDPM, FCNN
from DDPM.Config import *
from DEAD.AutoDecoder.Evaluation.InternalBallistic import vectoring_internal_ballistic
from DDPM.Evaluation.InternalBallistic import convert_av1_to_value, convert_av2_to_value
import shutil
from DDPM.Evaluation.Control import *

# 设置设备为GPU或CPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

def setup_ddpm(x_size, hidden_size, time_embed_size, b_condition_size, b_condition_embed_size, a_condition_size, a_condition_embed_size, n_T, drop_prob, postfix=""):
    """
    设置和初始化DDPM模型。
    
    参数:
    - x_size: 输入数据的大小
    - hidden_size: 隐藏层的大小
    - time_embed_size: 时间嵌入的大小
    - b_condition_size: 内弹道条件的大小
    - b_condition_embed_size: 内弹道嵌入的大小
    - a_condition_size: 属性条件的大小
    - a_condition_embed_size: 属性条件嵌入的大小
    - n_T: 时间步数
    - drop_prob: dropout概率
    
    返回:
    - 初始化并加载权重的DDPM模型
    """
    hidden_sizes = [hidden_size] * depth
    nn_model = FCNN(hidden_sizes=hidden_sizes,
                    x_size=x_size,
                    time_embed_size=time_embed_size,
                    b_condition_size=b_condition_size,
                    b_condition_embed_size=b_condition_embed_size,
                    a_condition_size=a_condition_size,
                    a_condition_embed_size=a_condition_embed_size)
    nn_model.eval()
    ddpm = DDPM(nn_model=nn_model,
                betas=(1e-4, 0.02),
                n_T=n_T,
                device=device,
                drop_prob=drop_prob).to(device)
    ddpm.eval()
    ddpm.load_state_dict(torch.load(f"{model_load_path}/DDPM{postfix}.pth", weights_only=True))
    return ddpm

def generate_samples(n_sample, b_condition, a_condition, guide_w, postfix=""):
    """
    生成样本数据。
    
    参数:
    - n_sample: 生成的样本数量
    - b_condition: 内弹道条件
    - a_condition: 属性条件
    - guide_w: 引导权重
    - postfix: 模型后缀
    
    返回:
    - 生成的样本数据
    """
    with torch.no_grad():
        x_size = 128 + 3 + 2
        a_condition_size = 3
        b_condition_size = 128
        ddpm = setup_ddpm(x_size,
                         hidden_size,
                         time_embed_size,
                         b_condition_size,
                         b_condition_embed_size,
                         a_condition_size,
                         a_condition_embed_size,
                         n_T,
                         drop_prob,
                         postfix)
        x_gen = ddpm.sample_single_cond(n_sample, x_size, guide_w, b_condition, a_condition)
        #x_gen= ddpm.sample_single_uncond(n_sample, x_size, b_condition, a_condition)
        return x_gen

def get_condition_from_index(trainingset_path, psfd_index, item_index):
    """
    从索引获取条件数据。
    
    参数:
    - trainingset_path: 训练集路径
    - psfd_index: PSFD索引
    - item_index: 项目索引
    
    返回:
    - b条件和a条件
    """
    npz_path = f'{trainingset_path}/{psfd_index}'
    npz_file = np.load(npz_path + ".npz")
    b_condition = torch.tensor(npz_file["bvs"][item_index], dtype=torch.float32)
    a_condition = torch.tensor(npz_file["avs2"][item_index], dtype=torch.float32)
    return b_condition, a_condition

def convert_attribute1_to_vector(r_ref, R):
    """
    将属性1转换为向量。
    
    参数:
    - r_ref: 推进剂燃烧（燃速）特性
    - R: 装药半径
    
    返回:
    - 转换后的向量
    """
    av1 = np.array([(r_ref - r_ref_min) / (r_ref_max - r_ref_min),  # 推进剂燃烧（燃速）特性
                    (R - R_min) / (R_max - R_min)]) * 0.1  # 装药半径
    return av1

def convert_attribute2_to_vector(cstar, rho, n, At):
    """
    将属性2转换为向量。
    
    参数:
    - cstar: 特征速度
    - rho: 密度
    - n: 压力指数
    - At: 喷口面积
    
    返回:
    - 转换后的向量
    """
    av2 = np.array([(cstar * rho - cstar_min * rho_min) / (cstar_max * rho_max - cstar_min * rho_min),  # 推进剂能量特性
                    (n - n_min) / (n_max - n_min),  # 推进剂压力特性
                    At * 100.0])
    return av2

def condition_from_user(t_array, p_array, cstar, rho, n, At):
    """
    从用户输入获取条件数据。
    
    参数:
    - t_array: 时间数组
    - p_array: 压力数组
    - cstar: 特征速度
    - rho: 密度
    - n: 压力指数
    - At: 喷口面积
    
    返回:
    - b条件和a条件
    """
    t_action, t_vector, p_vector = vectoring_internal_ballistic(t_array, p_array)
    b_condition = torch.tensor(np.stack([t_vector, p_vector * 1e-6], axis=1), dtype=torch.float32)
    a_condition = torch.tensor(convert_attribute2_to_vector(cstar, rho, n, At), dtype=torch.float32)
    return b_condition, a_condition

def start_generation(control_params):
    print(f"##### Start generation of {control_params['project_name']} #####")
    # 固定所有随机种子
    seed=control_params["seed"]
    torch.manual_seed(seed)
    if device == "cuda":
        torch.cuda.manual_seed_all(seed)
    
    if control_params["user_defined"]:
        t_array_target=control_params["t_array_target"]
        p_array_target=control_params["p_array_target"]
        cstar=control_params["cstar"]
        rho=control_params["rho"]
        n=control_params["n"]
        At=control_params["At"]
        b_condition, a_condition = condition_from_user(t_array_target,p_array_target,cstar,rho,n,At)
    else:
        i_from_psfd=control_params["i_from_psfd"]
        b_condition, a_condition = get_condition_from_index(trainingset_path, i_from_psfd, 0)
        cstar_rho, n, At = convert_av2_to_value(a_condition)

    # 生成样本
    out = generate_samples(n_sample=1000, b_condition=b_condition, a_condition=a_condition, guide_w=control_params["guid_w"],postfix=control_params["postfix"])
    out = out.cpu().numpy()
    lv = out[:, :-5]
    gv = out[:, -5:-2]
    av1 = out[:, -2:]

    r_ref_list = []
    R_list = []
    for av1_item in av1:
        r_ref, R = convert_av1_to_value(av1_item)
        r_ref_list.append(r_ref)
        R_list.append(R)

    # 保存生成的样本
    project_path=f"{decode_result_save_path}/{control_params['project_name']}"
    # 删除文件夹（如果存在）
    if os.path.exists(project_path):
        shutil.rmtree(project_path)
    print(f"项目保存目录：{project_path}")
    # 创建新的文件夹
    os.makedirs(project_path, exist_ok=True)
    if control_params["user_defined"]:
        np.savez_compressed(f"{project_path}/Samples.npz",
                            lv=lv,
                            gv=gv,
                            av1=av1,
                            r_ref=r_ref_list,
                            R=R_list,
                            cstar=cstar,
                            rho=rho,
                            n=n,
                            At=At,
                            t_array_target=t_array_target,
                            p_array_target=p_array_target)
    else:
        np.savez_compressed(f"{project_path}/Samples.npz",
                            lv=lv,
                            gv=gv,
                            av1=av1,
                            r_ref=r_ref_list,
                            R=R_list,
                            cstar_rho=cstar_rho,
                            n=n,
                            At=At,
                            i_from_psfd=i_from_psfd)

if __name__ == "__main__":
    control_params=get_control_params(question_name="K正弦状燃面装药",
                 postfix="_epoch50000",
                 cluster_num=5,
                 guid_w=0.5)
    start_generation(control_params=control_params)